"""
将JSON数据导入SQLite数据库
数据库文件存放在当前目录（bazi-test目录）
"""
import json
import sqlite3
from pathlib import Path
from typing import List, Dict
import sys


def create_database(db_path: str):
    """
    创建数据库表结构
    
    Args:
        db_path: 数据库文件路径
    """
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    # 创建数据表
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS paragraphs (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            book_code TEXT,
            book_name TEXT,
            category TEXT,
            page_num INTEGER,
            page_id TEXT,
            title TEXT,
            heaven_stem TEXT,
            earth_branch TEXT,
            title_keywords TEXT,
            pid TEXT,
            order_num INTEGER,
            original_text TEXT,
            translated_text TEXT,
            bzcases TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    
    # 创建索引以提升查询性能
    cursor.execute('''
        CREATE INDEX IF NOT EXISTS idx_book_code ON paragraphs(book_code)
    ''')
    cursor.execute('''
        CREATE INDEX IF NOT EXISTS idx_book_name ON paragraphs(book_name)
    ''')
    cursor.execute('''
        CREATE INDEX IF NOT EXISTS idx_page_id ON paragraphs(page_id)
    ''')
    cursor.execute('''
        CREATE INDEX IF NOT EXISTS idx_title_keywords ON paragraphs(title_keywords)
    ''')
    cursor.execute('''
        CREATE INDEX IF NOT EXISTS idx_heaven_stem ON paragraphs(heaven_stem)
    ''')
    cursor.execute('''
        CREATE INDEX IF NOT EXISTS idx_earth_branch ON paragraphs(earth_branch)
    ''')
    
    conn.commit()
    conn.close()
    print(f"✅ 数据库表结构已创建: {db_path}")


def import_json_to_db(json_path: str, db_path: str):
    """
    将JSON文件数据导入SQLite数据库
    
    Args:
        json_path: JSON文件路径
        db_path: 数据库文件路径
    """
    # 检查JSON文件是否存在
    json_file = Path(json_path)
    if not json_file.exists():
        print(f"❌ JSON文件不存在: {json_path}")
        sys.exit(1)
    
    # 读取JSON数据
    print(f"📄 正在读取JSON文件: {json_path}")
    with open(json_file, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    if not isinstance(data, list):
        print(f"❌ JSON文件格式错误: 应该是数组格式")
        sys.exit(1)
    
    print(f"📊 共 {len(data)} 条记录需要导入")
    
    # 连接数据库
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    # 准备插入SQL
    insert_sql = '''
        INSERT INTO paragraphs (
            book_code, book_name, category, page_num, page_id,
            title, heaven_stem, earth_branch, title_keywords,
            pid, order_num, original_text, translated_text, bzcases
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
    '''
    
    # 批量插入数据
    success_count = 0
    error_count = 0
    
    for idx, record in enumerate(data, 1):
        try:
            # 处理bzcases字段：如果是列表或字典，转换为JSON字符串；如果是None，保持None
            bzcases_str = None
            if record.get('bzcases'):
                bzcases_str = json.dumps(record.get('bzcases'), ensure_ascii=False)
            
            # 插入数据
            cursor.execute(insert_sql, (
                record.get('book_code'),
                record.get('book_name'),
                record.get('category'),
                record.get('page_num'),
                record.get('page_id'),
                record.get('title'),
                record.get('heaven_stem'),
                record.get('earth_branch'),
                record.get('title_keywords'),
                record.get('pid'),
                record.get('order_num'),
                record.get('original_text'),
                record.get('translated_text'),
                bzcases_str
            ))
            success_count += 1
            
            # 每1000条记录提交一次，提升性能
            if idx % 1000 == 0:
                conn.commit()
                print(f"  ✅ 已导入 {idx}/{len(data)} 条记录...")
        
        except Exception as e:
            error_count += 1
            print(f"  ❌ 第 {idx} 条记录导入失败: {e}")
            if error_count > 10:
                print(f"  ⚠️  错误过多，停止导入")
                break
    
    # 提交剩余的数据
    conn.commit()
    
    # 获取导入后的记录总数
    cursor.execute('SELECT COUNT(*) FROM paragraphs')
    total_count = cursor.fetchone()[0]
    
    conn.close()
    
    print(f"\n✅ 导入完成:")
    print(f"  成功导入: {success_count} 条记录")
    print(f"  导入失败: {error_count} 条记录")
    print(f"  数据库总记录数: {total_count} 条")


def main():
    """主函数"""
    # 数据库文件路径（放在当前目录）
    script_dir = Path(__file__).parent
    db_path = script_dir / "bazi_data.db"
    
    # JSON文件路径（默认使用data-analysis目录下的all_data.json）
    if len(sys.argv) > 1:
        json_path = sys.argv[1]
    else:
        json_path = script_dir / "data-analysis" / "all_data.json"
    
    # 创建数据库表结构
    create_database(str(db_path))
    
    # 导入数据
    import_json_to_db(str(json_path), str(db_path))
    
    print(f"\n📦 数据库文件位置: {db_path}")


if __name__ == '__main__':
    main()

